﻿import  os
import pandas as pd
import jieba
import re
import json
from elasticsearch import Elasticsearch
import json


#自然语句转换成logical Form
class Search:
    def __init__(self):
        dir = os.getcwd()
        self.QA_df=dir+r'/问答三元组.txt'
        self.QA_df=pd.read_csv(self.QA_df,encoding='utf-8',delimiter=' ')
        self.attr_dict={}
        self.attr=pd.read_csv('attr_mapping(1).txt',encoding='utf-8',header=None,index_col=None,delimiter='\t')
        self.attr_dict=self.attr.set_index(0).T.to_dict('list')

        #模糊匹配
        for i in self.attr_dict.keys():
            self.attr_dict[i] = [attr for attr in self.attr_dict[i] if str(attr) != 'nan']

        self.entity_set=set(self.QA_df['S(subject)'].tolist())
        self.rel_set=set(self.QA_df['P(predicate)'].tolist())
        self.attr_value_set=set(self.QA_df['O(object)'].tolist())

        # 将实体名、关系名、同义词、属性值加入到分词词典中
        for w in self.entity_set:
            jieba.add_word(w)
        for w in self.rel_set:
            jieba.add_word(w)
        self.attr_list=re.sub("[\[\]{}\' ]","",str(self.attr_dict)).replace(':',',').split(',')
        for w in self.attr_list:
            jieba.add_word(w)
        for w in self.attr_value_set:
            jieba.add_word(w)
        for w in ['不大于','不多于','不少于','不小于','不超过','','','',]:
            jieba.add_word(w)


    #自然语句转换成logical Form
    def ChatBot(self, sentence):
        q_sep=list(jieba.cut(sentence))
        # 识别属性名
        rel_match=[]
        es = Elasticsearch()
        for w in q_sep:
            if w in self.rel_set:
                rel_match.extend([w])
            rel_match.extend([k for k, v in self.attr_dict.items() if w in v])
        # 识别实体名
        entity_match=[]
        for w in q_sep:
            if w in self.entity_set:
                entity_match.append(w)
        # 识别属性值
        attr_value_match=[]
        for w in q_sep:
            if w in self.attr_value_set:
                attr_value_match.append(w)
        lf=None
        # logical form
        # 实体检索 e.g.:新冠病毒
        if len(rel_match)==0: # 如果没有属性名，则查询的是实体
            query = []
            lf=entity_match[0]
            query = json.dumps({"query": { "bool":{"filter":{"term" :{"subj" : lf}}}}},ensure_ascii=False)
            result = es.search(index='demo', body=query)
            #解析Json
            ans = ''
            for e in result['hits']['hits']:
                subj = e['_source']['subj']
                pred = e['_source']['po'][0]['pred']
                name = subj + pred
                ans = ans + name + ":" + e['_source']['po'][0]['obj'] + '\n'
            ans = ans[:-1]
        # 实体的属性检索 e.g.:新冠病毒的定义
        # 实体属性的多跳检索 e.g.:广东的省会的感染人数
        elif len(entity_match)!=0: # 如果含有属性名和实体名，则是属性查询
            lf=entity_match[0]+':'+':'.join(rel_match)
            divide = lf.replace(" ","").split(':')
            for i in range(len(divide)-1):
                if i == 0:
                    query = []
                    entity = divide[i]
                    pred = divide[i+1]
                    query.append(json.dumps({"query": { "bool":{"must":[{"match":  {"subj": entity}},
                                            {"nested":{"path":"po","query":{"bool":{"must":[{"term": { "po.pred": pred }}]}}}}]}}},ensure_ascii=False))
                    result = es.search(index='demo', body=query[0])
                    ans = result['hits']['hits'][0]['_source']['po'][0]['obj']
                if i == 1:
                    pred = divide[i+1]
                    query = []
                    query.append(json.dumps({"query": { "bool":{"must":[{"match":  {"subj": ans}},
                                            {"nested":{"path":"po","query":{"bool":{"must":[{"term": { "po.pred": pred }}]}}}}]}}},ensure_ascii=False))
                    result = es.search(index='demo', body=query[0])
                    ans = result['hits']['hits'][0]['_source']['po'][0]['obj']
        # 多种属性条件检索实体 e.g.:累计确诊人数超过1000人的中国省份或美国州
        # 如果含有属性名、不含实体名，则是根据属性值检索实体
        elif len(entity_match)==0:
            op=None
            if '超过' in q_sep or '大于' in q_sep or '多于' in q_sep or '以上' in q_sep:
                op='>'
            elif '小于' in q_sep or '少于' in q_sep or '以下' in q_sep:
                op='<'
            elif '不超过' in q_sep or '不大于' in q_sep or '不多于' in q_sep:
                op='<='
            elif '不小于' in q_sep or '不少于' in q_sep:
                op='>='
            num=re.findall(r"\d+",sentence)[0]
            if len(attr_value_match)!=0:
                attr_value_match=['所属:'+i for i in attr_value_match]
                attr_value_match[1:]=['OR'+i for i in attr_value_match[1:]]
                lf=rel_match[0]+op+num+'AND'+''.join(attr_value_match)
            ans_list = []
            po_list = []
            cmp_dir = {
                "<":"lt",
                "<=":"lte",
                ">":"gt",
                ">=":"gte"
            }

            divide = re.split('AND|OR', lf.replace(" ", ""))
            bool_ops = re.findall('AND|OR', lf)

            for e in divide:
                if e == "":
                    return "", 'AND 或 OR 后不能为空'

                begin_with_NOT = False
                if e[0:3] == 'NOT':
                    begin_with_NOT = True
                    e = e[3:]
                elif 'NOT' in e:
                    ans = 'NOT请放在PO对前面'

                op = re.findall("：|:|>|<|>=|<=",e)
                if len(op) != 1:
                    ans = '语法错误'
                op = op[0]

                if op == '<' or op == '>':
                    index = e.find(op)
                    if e[index+1] == '=':
                        op = op + '='
                pred, obj = e.split(op)

                
                part_query = ""
                if not begin_with_NOT:
                    if op == ':' or op == '：':
                        if pred == 'subj' or pred == '累计确诊':
                            part_query = '{"term":{"' + pred + '":' + obj + '}}'
                        else:
                            part_query = '{"nested":{"path":"po","query":{"bool":{"must":[{"term":{"po.pred":"' + pred + \
                                    '"}},{"term":{"po.obj":"' + obj + '"}}]}}}}'
                    else:
                        if pred == '累计确诊':
                            part_query = '{"range":{"' + pred + '":{"' + cmp_dir[op] + '":' + obj + '}}}'
                        else:
                            ans = '该属性不支持比较大小,目前只支持累计确诊'
                else:
                    if op == ':' or op == '：':
                        if pred == 'subj' or pred == '累计确诊':
                            part_query = '{"bool":{"must_not":{"term":{"' + pred + '":' + obj + '}}}}'
                        else:
                            part_query = '{"nested":{"path":"po","query":{"bool":{"must":[{"term":{"po.pred":"' + pred + \
                                    '"}},{"bool":{"must_not":{"term":{"po.obj":"' + obj + '"}}}}]}}}}'
                    else:
                        if pred == '累计确诊':
                            part_query = '{"bool":{"must_not":{"range":{"' + pred + '":{"' + cmp_dir[op] + '":' + obj + \
                                    '}}}}}'
                        else:
                            ans = '该属性不支持比较大小,目前只支持累计确诊'        
                po_list.append(part_query)

            or_po = [False] * len(divide)
            should_list = []
            must_list = []
            i = 0
            while i < len(bool_ops):
                if bool_ops[i] == 'OR':
                    adjacent_or = [po_list[i]]
                    or_po[i] = True
                    while i < len(bool_ops) and bool_ops[i] == 'OR':
                        adjacent_or.append(po_list[i+1])
                        or_po[i+1] = True
                        i += 1
                    should_list.append(",".join(adjacent_or))
                i += 1
            for i,po in enumerate(or_po):
                if not po:
                    must_list.append(po_list[i])
            must_list = ",".join(must_list)
            query = ""
            if must_list:
                query = '{"query":{"bool":{"must":[' + must_list + ']'
                if should_list:
                    query += ","
                    for s in should_list:
                        query += '"should":[' + s + '],'
                    query = query[:-1]
                query += '}}}' 
            else:
                query = '{"query":{"bool":{'
                if should_list:
                    for s in should_list:
                        query += '"should":[' + s + '],'
                    query = query[:-1]
                query += '}}}'
            result = es.search(index='demo', body=query,size=100)
            ans = []
            for e in result['hits']['hits']:
                subj = e['_source']['subj']
                rel = rel_match[0]
                attr_value = e['_source'][rel]
                value = subj + "：" + rel + str(attr_value)
                ans.append(value)
            ans = '\n'.join(ans)
        return ans